// clang-format off
//
// Created by goksu on 4/6/19.
//

#include <algorithm>
#include <vector>
#include "rasterizer.hpp"
#include <opencv2/opencv.hpp>
#include <array>
#include <math.h>
#include<algorithm>


rst::pos_buf_id rst::rasterizer::load_positions(const std::vector<Eigen::Vector3f> &positions) {
  auto id = get_next_id();
  pos_buf.emplace(id, positions);

  return {id};
}

rst::ind_buf_id rst::rasterizer::load_indices(const std::vector<Eigen::Vector3i> &indices) {
  auto id = get_next_id();
  ind_buf.emplace(id, indices);

  return {id};
}

rst::col_buf_id rst::rasterizer::load_colors(const std::vector<Eigen::Vector3f> &cols) {
  auto id = get_next_id();
  col_buf.emplace(id, cols);

  return {id};
}

auto to_vec4(const Eigen::Vector3f &v3, float w = 1.0f) {
  return Vector4f(v3.x(), v3.y(), v3.z(), w);
}

std::vector<Vector2f> rst::rasterizer::subPixelPostions(int x, int y){
  std::vector<Vector2f> samplingPositions;
  float step = 1.0f / ((float)msaa / 2.0);
  for (int r = 0; r < msaa / 2; ++r) {
    for (int c = 0; c < msaa / 2; ++c) {
      auto p = Vector2f {
          (float)x + c * step + 0.5 * step,
          (float)y + r * step + 0.5 * step,
      };
      samplingPositions.push_back(p);
    }
  }
  return samplingPositions;
}

static bool insideTriangle(float x, float y, const Vector3f *_v) {
  // a b c , _v[0], _v[1], _v[2]
  Vector3f ab = Vector3f{_v[1].x() - _v[0].x(), _v[1].y() - _v[0].y(), 0};
  Vector3f ap{x - _v[0].x(), y - _v[0].y(), 0};

  Vector3f bc = Vector3f{_v[2].x() - _v[1].x(), _v[2].y() - _v[1].y(), 0};
  Vector3f bp{x - _v[1].x(), y - _v[1].y(), 0};

  Vector3f ca = Vector3f{_v[0].x() - _v[2].x(), _v[0].y() - _v[2].y(), 0};
  Vector3f cp{x - _v[2].x(), y - _v[2].y(), 0};
  return (ab.cross(ap).z() > 0
          && bc.cross(bp).z() > 0
          && ca.cross(cp).z() > 0)
         ||
         (ab.cross(ap).z() < 0
          && bc.cross(bp).z() < 0
          && ca.cross(cp).z() < 0);
}

static std::tuple<float, float, float> computeBarycentric2D(float x, float y, const Vector3f *v) {
  float c1 = (x * (v[1].y() - v[2].y()) + (v[2].x() - v[1].x()) * y + v[1].x() * v[2].y() -
              v[2].x() * v[1].y()) /
             (v[0].x() * (v[1].y() - v[2].y()) + (v[2].x() - v[1].x()) * v[0].y() +
              v[1].x() * v[2].y() - v[2].x() * v[1].y());
  float c2 = (x * (v[2].y() - v[0].y()) + (v[0].x() - v[2].x()) * y + v[2].x() * v[0].y() -
              v[0].x() * v[2].y()) /
             (v[1].x() * (v[2].y() - v[0].y()) + (v[0].x() - v[2].x()) * v[1].y() +
              v[2].x() * v[0].y() - v[0].x() * v[2].y());
  float c3 = (x * (v[0].y() - v[1].y()) + (v[1].x() - v[0].x()) * y + v[0].x() * v[1].y() -
              v[1].x() * v[0].y()) /
             (v[2].x() * (v[0].y() - v[1].y()) + (v[1].x() - v[0].x()) * v[2].y() +
              v[0].x() * v[1].y() - v[1].x() * v[0].y());
  return {c1, c2, c3};
}

void rst::rasterizer::draw(pos_buf_id pos_buffer, ind_buf_id ind_buffer, col_buf_id col_buffer,
                           Primitive type) {
  auto &buf = pos_buf[pos_buffer.pos_id];
  auto &ind = ind_buf[ind_buffer.ind_id];
  auto &col = col_buf[col_buffer.col_id];

  float f1 = (50 - 0.1) / 2.0;
  float f2 = (50 + 0.1) / 2.0;

  Eigen::Matrix4f mvp = projection * view * model;
  for (auto &i : ind) {
    Triangle t;
    Eigen::Vector4f v[] = {
        mvp * to_vec4(buf[i[0]], 1.0f),
        mvp * to_vec4(buf[i[1]], 1.0f),
        mvp * to_vec4(buf[i[2]], 1.0f)
    };
    //Homogeneous division
    for (auto &vec : v) {
      vec /= vec.w();
    }
    //Viewport transformation
    for (auto &vert : v) {
      vert.x() = 0.5 * width * (vert.x() + 1.0);
      vert.y() = 0.5 * height * (vert.y() + 1.0);
      vert.z() = vert.z() * f1 + f2;
    }

    for (int i = 0; i < 3; ++i) {
      t.setVertex(i, v[i].head<3>());
    }

    auto col_x = col[i[0]];
    auto col_y = col[i[1]];
    auto col_z = col[i[2]];

    t.setColor(0, col_x[0], col_x[1], col_x[2]);
    t.setColor(1, col_y[0], col_y[1], col_y[2]);
    t.setColor(2, col_z[0], col_z[1], col_z[2]);

    if (msaa == 1) {
      rasterize_triangle_for_no_msaa(t);
    } else {
      rasterize_triangle_for_msaa(t);
    }

  }
}

// {min,max}
std::array<Vector2i, 2> getAabb(const Triangle &t) {
  int minX = ceil(fmin(fmin(t.v[0].x(), t.v[1].x()), t.v[2].x()));
  int minY = ceil(fmin(fmin(t.v[0].y(), t.v[1].y()), t.v[2].y()));
  int maxX = floor(fmax(fmax(t.v[0].x(), t.v[1].x()), t.v[2].x()));
  int maxY = floor(fmax(fmax(t.v[0].y(), t.v[1].y()), t.v[2].y()));
  return {
      Vector2i{minX, minY},
      Vector2i{maxX, maxY},
  };
}

void rst::rasterizer::rasterize_triangle_for_no_msaa(const Triangle &t) {
  std::array<Vector4f, 3> v = t.toVector4();
  std::array<Vector2i, 2> aabb = getAabb(t);
  //遍历 (minX,minY) -> (maxX,maxY) 之间的所有点
  int minX = aabb[0].x();
  int minY = aabb[0].y();
  int maxX = aabb[1].x();
  int maxY = aabb[1].y();
  for (int x = minX; x <= maxX; ++x) {
    for (int y = minY; y <= maxY; ++y) {
      if (insideTriangle(x, y, t.v)) {
        // z插值
        auto[alpha, beta, gamma] = computeBarycentric2D(x, y, t.v);
        float w_reciprocal = 1.0 / (alpha / v[0].w() + beta / v[1].w() + gamma / v[2].w());
        float z_interpolated = alpha * v[0].z() / v[0].w() + beta * v[1].z() / v[1].w() +
                               gamma * v[2].z() / v[2].w();
        z_interpolated *= w_reciprocal;
        if (z_interpolated < depth_buf[get_index(x, y)]) {
          set_pixel(Vector3f {x, y, z_interpolated}, t.getColor());
          depth_buf[get_index(x, y)] = z_interpolated;
          frame_buf[get_index(x, y)] = t.getColor();
        }
      }
    }
  }
}
//Screen space rasterization
void rst::rasterizer::rasterize_triangle_for_msaa(const Triangle &t) {
  std::array<Vector4f, 3> v = t.toVector4();
  std::array<Vector2i, 2> aabb = getAabb(t);
  //遍历 (minX,minY) -> (maxX,maxY) 之间的所有点
  int minX = aabb[0].x();
  int minY = aabb[0].y();
  int maxX = aabb[1].x();
  int maxY = aabb[1].y();
  for (int x = minX; x <= maxX; ++x) {
    for (int y = minY; y <= maxY; ++y) {
      // 先使用子像素绘制到高分辨率的FBO
      auto subPixelPositions = subPixelPostions(x, y);
      for (int subPixelIndex = 0; subPixelIndex < subPixelPositions.size(); ++subPixelIndex) {
        Vector2f subPixelPostion = subPixelPositions[subPixelIndex];
        if (insideTriangle(subPixelPostion.x(), subPixelPostion.y(), t.v)) {
          // z插值
          auto[alpha, beta, gamma] = computeBarycentric2D(subPixelPostion.x(), subPixelPostion.y(),
                                                          t.v);
          float w_reciprocal = 1.0 / (alpha / v[0].w() + beta / v[1].w() + gamma / v[2].w());
          float z_interpolated = alpha * v[0].z() / v[0].w() + beta * v[1].z() / v[1].w() +
                                 gamma * v[2].z() / v[2].w();
          z_interpolated *= w_reciprocal;
          if (z_interpolated < depth_buf_msaa[get_msaa_buffer_index(x, y, subPixelIndex)]) {
            set_pixel_msaa(x, y, subPixelIndex, t.getColor());
            depth_buf_msaa[get_msaa_buffer_index(x, y, subPixelIndex)] = z_interpolated;
          }
        }
      }
      // 将高分辨率FBO上的多个点求平均值得到一个像素点的值
      // down sampling
      Vector3f totalColor = {0,0,0};
      float totalDepth = 0;
      for (int subPixelIndex = 0; subPixelIndex < msaa; ++subPixelIndex) {
        int idx = get_msaa_buffer_index(x, y, subPixelIndex);
        totalColor += frame_buf_msaa[idx];
        totalDepth += depth_buf_msaa[idx];
      }
      Vector3f avgColor = totalColor / (float)msaa;
      float avgDepth = totalDepth / (float)msaa;
      set_pixel(Vector3f{x, y, avgDepth}, avgColor);
      depth_buf[get_index(x, y)] = avgDepth;
    }
  }
}

void rst::rasterizer::set_model(const Eigen::Matrix4f &m) {
  model = m;
}

void rst::rasterizer::set_view(const Eigen::Matrix4f &v) {
  view = v;
}

void rst::rasterizer::set_projection(const Eigen::Matrix4f &p) {
  projection = p;
}

void rst::rasterizer::clear(rst::Buffers buff) {
  if ((buff & rst::Buffers::Color) == rst::Buffers::Color) {
    std::fill(frame_buf.begin(), frame_buf.end(), Eigen::Vector3f{0, 0, 0});
    std::fill(frame_buf_msaa.begin(), frame_buf_msaa.end(), Eigen::Vector3f{0, 0, 0});
  }
  if ((buff & rst::Buffers::Depth) == rst::Buffers::Depth) {
    std::fill(depth_buf.begin(), depth_buf.end(), std::numeric_limits<float>::infinity());
    std::fill(depth_buf_msaa.begin(), depth_buf_msaa.end(), std::numeric_limits<float>::infinity());
  }
}

rst::rasterizer::rasterizer(int w, int h) : width(w), height(h) {
  frame_buf.resize(w * h);
  depth_buf.resize(w * h);
  frame_buf_msaa.resize(width * msaa * height);
  depth_buf_msaa.resize(width * msaa * height);
}

int rst::rasterizer::get_index(int x, int y) {
  return (height - 1 - y) * width + x;
}

int rst::rasterizer::get_msaa_buffer_index(int x, int y, int subPixelIndex)
{
  return y * width * msaa + x * msaa + subPixelIndex;
}


void rst::rasterizer::set_pixel(const Eigen::Vector3f &point, const Eigen::Vector3f &color) {

  //old index: auto ind = point.y() + point.x() * width;
  auto ind = (height - 1 - point.y()) * width + point.x();
  frame_buf[ind] = color;
}

void rst::rasterizer::set_pixel_msaa(int x, int y, int subPixelIndex, const Eigen::Vector3f &color) {
  frame_buf_msaa[get_msaa_buffer_index(x, y, subPixelIndex)] = color;
}

// clang-format on